package com.lw.leetcode.tree.b;

import java.util.ArrayList;
import java.util.List;

/**
 * Created with IntelliJ IDEA.
 * 1993. 树上的操作
 *
 * @author liw
 * @version 1.0
 * @date 2022/5/25 9:29
 */
public class LockingTree {

    private int n;
    private int[] locked;
    private List<Integer>[] children;
    private int[] parent;

    public LockingTree(int[] parent) {
        this.n = parent.length;
        this.locked = new int[n];
        this.children = new ArrayList[n];
        this.parent = parent;
        for (int i = 0; i < n; i++) {
            children[i] = new ArrayList<>();
        }
        for (int i = 0; i < n; i++) {
            if (parent[i] != -1) {
                children[parent[i]].add(i);
            }
        }
    }

    public boolean lock(int num, int user) {
        if (locked[num] == 0) {
            locked[num] = user;
            return true;
        }
        return false;
    }

    public boolean unlock(int num, int user) {
        if (locked[num] == user) {
            locked[num] = 0;
            return true;
        }
        return false;
    }

    public boolean upgrade(int num, int user) {
        if (locked[num] == 0 && this.checkUnlock(num) && this.checkLock(num)) {
            locked[num] = user;
            this.unlock(num);
            return true;
        }
        return false;
    }

    private void unlock(int num) {
        List<Integer> cur = children[num];
        for (int child : cur) {
            if (this.locked[child] != 0) {
                this.locked[child] = 0;
            }
            this.unlock(child);
        }
    }

    private boolean checkLock(int num) {
        List<Integer> cur = children[num];
        for (int child : cur) {
            if (this.locked[child] != 0) {
                return true;
            }
            if (this.checkLock(child)) {
                return true;
            }
        }
        return false;
    }

    private boolean checkUnlock(int num) {
        while (parent[num] != -1 && this.locked[parent[num]] == 0) {
            num = parent[num];
        }
        return parent[num] == -1;
    }

}
